package org.iscas.tcse.datasample;

import org.apache.commons.math3.linear.MatrixUtils;
import org.apache.commons.math3.linear.RealVector;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.ml.linalg.DenseVector;
import org.apache.flink.streaming.api.functions.ProcessFunction;
import org.apache.flink.util.Collector;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import redis.clients.jedis.Jedis;

import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.concurrent.atomic.AtomicInteger;

public class CoresetSampler extends ProcessFunction<DenseVector, DenseVector> implements Sampler {
    private static final Logger LOG = LoggerFactory.getLogger(CoresetSampler.class);
    private final int numCategories;

    private final int windowSize;
    //    private final int perCategorySelectCnt;

    private volatile double sampleRatio;
    private volatile AtomicInteger counter;
    private final int numFeatures;

    private int buffLen;

    private RealVector[] samples;

    private double[][] distanceMatrix;
    private int[] labels;

    private double[] rowBuffer;

    private final SamplerConf conf;

    private SampleController controller = null;
    private final boolean enableController;

    private volatile boolean isRunning = true;

    public CoresetSampler(
            int numFeatures,
            int windowSize,
            double sampleRatio,
            int numCategories,
            SamplerConf conf,
            boolean enableController) {
        this.numCategories = numCategories;
        this.windowSize = windowSize;
        this.sampleRatio = sampleRatio;
        this.numFeatures = numFeatures;
        this.distanceMatrix = new double[windowSize][windowSize];
        this.labels = new int[windowSize];
        this.conf = conf;
        this.enableController = enableController;
    }

    @Override
    public void open(Configuration parameters) throws Exception {
        super.open(parameters);
        this.samples = new RealVector[windowSize];
        this.labels = new int[windowSize];
        this.rowBuffer = new double[numFeatures];
        this.buffLen = 0;

        if (enableController) {
            int taskId = getRuntimeContext().getIndexOfThisSubtask();
            if (taskId == 0) {
                Thread controller = new Thread(new SampleController(this.conf));
                controller.start();
                this.controller = (SampleController) controller;
                this.counter = new AtomicInteger(0);
            }

            new Thread(
                            () -> {
                                try {
                                    syncSampleRate();
                                } catch (InterruptedException e) {
                                    e.printStackTrace();
                                }
                            })
                    .start();
        }
    }

    @Override
    public void close() throws Exception {
        super.close();
        if (this.controller != null) {
            this.controller.stopRunning();
        }
        this.isRunning = false;
    }

    double deltaB(Set<Integer> selected, List<Integer> allSamples) {
        double allDis = 0.0;
        for (int i : allSamples) {
            double minDis = Double.MAX_VALUE;
            if (selected.contains(i)) {
                continue;
            }
            for (int j : selected) {
                minDis = Math.min(minDis, distanceMatrix[i][j]);
            }
            allDis += minDis;
        }
        return allDis;
    }

    void calDistanceMatrix() {
        for (int i = 0; i < windowSize; i++) {
            for (int j = 0; j < i; j++) {
                distanceMatrix[i][j] = distanceMatrix[j][i];
            }
            for (int j = i + 1; j < windowSize; j++) {
                distanceMatrix[i][j] = samples[i].getDistance(samples[j]);
            }
        }
    }

    Set<Integer> submodularMaximization() {
        calDistanceMatrix();
        Set<Integer> selected = new HashSet<>();
        //        int perCategorySelectCnt = (int) (sampleRatio * windowSize / numCategories);
        int[] perCategorySelectCnts = new int[numCategories];
        for (int i = 0; i < windowSize; i++) {
            perCategorySelectCnts[labels[i]]++;
        }
        for (int i = 0; i < numCategories; i++) {
            perCategorySelectCnts[i] = (int) Math.round(perCategorySelectCnts[i] * sampleRatio);
        }

        for (int u = 0; u < numCategories; u++) {
            int perCategorySelectCnt = perCategorySelectCnts[u];

            List<Integer> indicesOfU = new ArrayList<>();
            Set<Integer> currentSelected = new HashSet<>();
            for (int i = 0; i < windowSize; i++) {
                if (labels[i] == u) {
                    indicesOfU.add(i);
                }
            }
            // 数量不足，全部选中
            if (indicesOfU.size() <= perCategorySelectCnt) {
                selected.addAll(indicesOfU);
                continue;
            }
            for (int i = 0; i < perCategorySelectCnt; i++) {
                double minDis = Double.MAX_VALUE;
                int select = 0;
                for (int index : indicesOfU) {
                    if (!currentSelected.contains(index)) {
                        currentSelected.add(index);
                        double allDis = deltaB(currentSelected, indicesOfU);
                        if (allDis < minDis) {
                            minDis = allDis;
                            select = index;
                        }
                        currentSelected.remove(index);
                    }
                }
                //                LOG.info("select {}, diss {}", select, minDis);
                currentSelected.add(select);
            }
            selected.addAll(currentSelected);
        }
        return selected;
    }

    private DenseVector constructVector(int idx) {
        int len = numFeatures + 1;
        double[] values = new double[len];
        for (int i = 0; i < numFeatures; i++) {
            values[i] = samples[idx].getEntry(i);
        }
        values[numFeatures] = labels[idx];
        return new DenseVector(values);
    }

    @Override
    public void processElement(
            DenseVector vector,
            ProcessFunction<DenseVector, DenseVector>.Context ctx,
            Collector<DenseVector> out)
            throws Exception {
        for (int i = 0; i < numFeatures; i++) {
            rowBuffer[i] = vector.get(i);
        }
        labels[buffLen] = (int) vector.get(numFeatures);
        samples[buffLen] = MatrixUtils.createRealVector(rowBuffer);
        buffLen += 1;
        if (buffLen >= windowSize) {
            buffLen = 0;
            //            LOG.info("Start to select samples");
            Set<Integer> selected = submodularMaximization();
            //            LOG.info("Selected {} samples", selected.size());
            for (int idx : selected) {
                out.collect(constructVector(idx));
            }
        }
        counter.getAndIncrement();
    }

    public void syncSampleRate() throws InterruptedException {
        Jedis jedis = new Jedis(conf.redisHost, conf.redisPort);
        while (isRunning) {
            String sampleRatioStr = jedis.get(conf.sampleRatioKey);
            if (sampleRatioStr != null) {
                this.sampleRatio = Double.parseDouble(sampleRatioStr);
            }
            Thread.sleep(conf.syncInterval);
        }
        jedis.close();
    }

    @Override
    public void reportSourceSpeed() throws InterruptedException {
        Jedis jedis = new Jedis(conf.redisHost, conf.redisPort);
        while (isRunning) {
            jedis.set(conf.sourceSpeedKey, String.valueOf(counter));
            counter.set(0);
            Thread.sleep(conf.syncInterval);
        }
        jedis.close();
    }
}
